'''
算法题：使用高斯消元法求解线性方程组
给定一个线性方程组AX=B，其中A为nxn的方阵，B为nx1的列向量，使用高斯消元法求其根X。
高斯消元法：对增广矩阵[A,B]进行行变换（交换两行，给某一行乘以非0系数，将某一行乘以系数加在另一行上），最终变为[A',B']的形式，其中A'为上三角矩阵。再反解A'X=B'即可。
说明：若方程有唯一解，返回nx1的列向量，为方程的根。若方程无解或有多个解，返回False

举例：
    输入：A=[[1, 0, 0], [0, 1, 0], [0, 0, 1]]  B = [[1], [2], [3]]
    输出：[[1], [2], [3]]

    输入：A=[[1, 0, 0], [0, 1, 0], [0, 0, 0]]  B = [[0], [0], [1]]
    输出：False

'''
import numpy as np

def solve(A, B):
    # A is a np.array with size (n, n)
    # B is a np.array with size (n, 1)
    array_size = A.shape[0]

    #在这里实现
    # pass
    # concat
    Z = np.concatenate((A, B), axis=1)


    for j in range(array_size):
        # 换行操作
        if Z[j][j] == 0:
            if j == array_size-1:
                return False

            for change in range(j+1,array_size-1):
                if Z[change][j] != 0:
                    Z[j],Z[change] = Z[change],Z[j]
        # 消元
        for i in range(j,array_size-1):
            alpha_ = - (Z[i+1][j] / Z[j][j])

            Z[i+1] += Z[j]* alpha_

        Z[j] = Z[j] / Z[j][j]
    # 消元ok
    # return Z

    # 回代求解

    X = [0] * (array_size)

    if Z[array_size-1][array_size-1] == 0:
        # I 阵不满秩则解不唯一 False
        return False

    else:
        # 满秩则回代求解
        for x_i in range(array_size-1,-1,-1):
            if x_i == array_size-1:
                X[x_i] = Z[x_i][array_size]
            else:
                temp = 0
                # 这里用递归感觉写会更简洁
                for m in range(x_i+1,array_size):
                    temp += (Z[x_i][m] * X[m])

                X[x_i] = Z[x_i][array_size] - temp

        return np.array(X).reshape(-1,1)

if __name__ == '__main__':
    # A = np.array( [[1, 0, 0], [0, 1, 0], [0, 0, 0]] ,dtype= 'float64')
    # B = np.array( [[0], [0], [1]],dtype= 'float64')
    A = np.array( [[1, 0, 0], [0, 1, 0], [0, 0, 1]],dtype= 'float64')
    B = np.array( [[1], [2], [3]] ,dtype= 'float64')
    # A = np.array([[2, 23, 5, 1], [2, 5, 9, 12], [4, 2, 34, 2], [1, 5, 3, 61]],dtype= 'float64')
    # B = np.array([[59], [-16], [126], [-284]],dtype= 'float64')
    X = solve(A, B)
    print(X)
